// Copyright (c) 2024 by Rockchip Electronics Co., Ltd. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.


#ifndef _RKNN_SV2_YOLO11_H_
#define _RKNN_SV2_YOLO11_H_

#include "sms_core.h"
#include <string>
#include <chrono>
#include <mutex>
#include <queue>
#include <condition_variable>
#include <sys/utsname.h>
#include <atomic>
#include "yolo11.h"
namespace sv2
{
  //class rknnYolov11;
  class ImgMsg {
    public:
      ImgMsg(cv::Mat img_, nlohmann::json msg_) : img(img_), msg(msg_) {}
      cv::Mat img;
      nlohmann::json msg;
  };

  class YOLOv11DetCppNode_Rknn : public sms::BaseNode
  {
  public:
    YOLOv11DetCppNode_Rknn(
        std::string job_name,
        std::string param_file,
        std::string ip = "127.0.0.1",
        int port = 9094,
        nlohmann::json kwargs = {}
        ): sms::BaseNode("YOLOv11DetCppNode_Rknn", job_name, param_file, ip, port, kwargs)
    {
      // 读取节点参数
      
      this->launch_next_emit =  this->get_param("launch_next_emit", true);
      this->specified_input_topic = this->get_param("specified_input_topic", "");
      this->specified_output_topic = this->get_param("specified_output_topic", "");
      this->realtime_det = this->get_param("realtime_det", true);
      this->remote_ip = this->get_param("remote_ip", "127.0.0.1");
      this->remote_port = this->get_param("remote_port", 9094);
      this->confidence = this->get_param("confidence", 0.25);
      this->nms_thresh = this->get_param("nms_thresh", 0.45);
      this->dataset_name = this->get_param("dataset_name", "coco_detection");
      this->model_path = this->get_param("model_path", "/home/orangepi/Downloads/yolo11n.rknn");//"/home/orangepi/Downloads/coco_detection-yolo11n_i512_c80-20250522.rknn");
      this->target = this->get_param("target", "rk3588");
      this->use_shm = this->get_param("use_shm", -1);
      this->g_dataset_categories = this->get_param("/dataset_categories", {});
      this->params_help();
      nlohmann::json dataset_categories = this->g_dataset_categories[this->dataset_name];
      
      this->get_img_flag = 0;

      struct utsname sysbuf;
      uname(&sysbuf);
      printf("sys %s\n", sysbuf.sysname);
      this->b_use_shm = false;
      if(this->use_shm == 1 || (this->use_shm == -1 && strcmp(sysbuf.sysname, "Linux") == 0))
      {
        this->b_use_shm = true;
      }
      // self.co_helper

      std::string input_topic = "/" + job_name + "/sensor/image_raw";
      std::string output1_topic = "/" + job_name + "/detector/results";
      std::string output2_topic = "/" + job_name + "/detector/image_results";
      if (this->specified_input_topic.size() > 0)
        input_topic = this->specified_input_topic;
      if (this->specified_output_topic.size() > 0)
        output1_topic = this->specified_output_topic;
  
      this->_det_res_sub = new sms::Subscriber(input_topic, "std_msgs::Null", std::bind(&YOLOv11DetCppNode_Rknn::full_res_callback, this, std::placeholders::_1));
      this->_result_writer = new sms::Publisher(output1_topic, "spirecv_msgs::2DTargets", sms::QoS::BestEffort, this->remote_ip, this->remote_port);
      this->_show_writer = new sms::Publisher(output2_topic, this->b_use_shm ? "memory_msgs::RawImage" : "sensor_msgs::CompressedImage");

      if(this->launch_next_emit)
      {
        this->_next_writer = new sms::Publisher("/live/launch_next", "std_msgs::Boolean");
      }
      /*
      logger.info("_imgsz: [" + this->_imgsz[0].dump() + ", " + this->_imgsz[1].dump() + "]");
      logger.info("port: " + std::to_string(static_cast<int>(this->remote_port)));
      logger.info("model_path: " + this->model_path);
      logger.info("dataset_name: " + this->dataset_name);
      logger.info("confidence: " + std::to_string(static_cast<float>(this->confidence)));
      logger.info("nms_thresh: " + std::to_string(static_cast<float>(this->nms_thresh)));
      */
      std::string abs_model_path;
      if (sms::_startswith(this->model_path, "sms::"))
      {
        abs_model_path = sms::download_model("YOLOv11DetCppNode_Rknn", this->model_path);
        assert(abs_model_path.size() > 0);
      }
      else
      {
        abs_model_path = this->model_path;
      }
      bool ret;
      this->_vs.getconfig(this->_result_writer, this->_show_writer, this->_next_writer, this->b_use_shm, this->launch_next_emit);

      ret = this->_vs.setup(abs_model_path, this->confidence, this->nms_thresh, this->g_dataset_categories, this->dataset_name);
      if (ret)
      {
          std::cout<<"Setup is OK"<<std::endl;
      }
      else
      {
          std::cout<<"Setup is Bad"<<std::endl;
      }
      ret = this->_vs.isOpened();
      if(ret)
      {
          std::cout<<"yolo11 channel is Opened"<<std::endl;
      }
      else
      {
          std::cout<<"yolov11 channel is Closed"<<std::endl;
      }
    }
    ~YOLOv11DetCppNode_Rknn()
    {
      if (_det_res_sub != NULL)
        delete _det_res_sub;
    }
    void run();

    void full_res_callback(nlohmann::json msg)
    {
      if (this->realtime_det && this->get_img_flag != 0)
      {
        return;
      }
      cv::Mat img = sms::sms2cvimg(msg);
      // std::cout << "get: " << img.cols << ", " << img.rows << std::endl;
      ImgMsg imgmsg(img, msg);
      // 放入到阻塞队列中
      this->_full_queue.push(imgmsg);
      this->get_img_flag = 1;
      // 通知主线程取数据
      // this->_full_cv.notify_one();
    }

    std::string _image_topic;
    std::string _image_out_topic;

    bool launch_next_emit;
    std::string specified_input_topic;
    std::string specified_output_topic;
    bool realtime_det;
    std::string remote_ip;
    int remote_port;
    float confidence;
    float nms_thresh;
    std::string dataset_name;
    std::string model_path;
    std::string target;
    int use_shm; 
    bool b_use_shm;
    nlohmann::json g_dataset_categories;
    nlohmann::json _imgsz;
    std::atomic<int> get_img_flag;
    rknnYolov11 _vs;

    sms::Publisher* _result_writer;
    sms::Publisher* _show_writer;
    sms::Publisher* _next_writer;
private:
    // 订阅话题
    sms::Subscriber* _det_res_sub;

    // std::mutex _full_queue_mtx;
    std::queue<ImgMsg> _full_queue;
    // std::condition_variable _full_cv;
  };

}

#endif //_RKNN_SV2_YOLO11_H_